#include <iostream>
#include <pcl/io/pcd_io.h>
#include <pcl/registration/ndt.h>
#include <pcl/filters/approximate_voxel_grid.h>
#include <pcl/filters/filter.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <thread>
#include <chrono>
#include <X11/Xlib.h>

pcl::PointCloud<pcl::PointXYZ>::Ptr read_cloud_point(std::string const &file_path) {
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
    if (pcl::io::loadPCDFile<pcl::PointXYZ>(file_path, *cloud) == -1) {
        PCL_ERROR("Couldn't read the pcd file\n");
        return nullptr;
    }
    
    // 移除NaN点
    std::vector<int> indices;
    pcl::removeNaNFromPointCloud(*cloud, *cloud, indices);
    return cloud;
}

void visualizer(pcl::PointCloud<pcl::PointXYZ>::Ptr target_cloud, 
               pcl::PointCloud<pcl::PointXYZ>::Ptr output_cloud) {
    if (target_cloud->empty() || output_cloud->empty()) {
        std::cerr << "Error: Empty cloud passed to visualizer!" << std::endl;
        return;
    }

    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer_final(new pcl::visualization::PCLVisualizer("3D Viewer"));
    viewer_final->setBackgroundColor(0, 0, 0);

    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> target_color(target_cloud, 255, 0, 0);
    viewer_final->addPointCloud<pcl::PointXYZ>(target_cloud, target_color, "target cloud");
    viewer_final->setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 1, "target cloud");

    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> output_color(output_cloud, 0, 255, 0);
    viewer_final->addPointCloud<pcl::PointXYZ>(output_cloud, output_color, "output cloud");
    viewer_final->setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 1, "output cloud");

    viewer_final->addCoordinateSystem(1.0, "global");
    viewer_final->initCameraParameters();

    while (!viewer_final->wasStopped()) {
        viewer_final->spinOnce(100);
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

int main(int argc, char **argv) {
    XInitThreads();  // 初始化X11多线程支持

    // 加载并检查点云
    auto target_cloud = read_cloud_point(argv[1]);
    auto input_cloud = read_cloud_point(argv[2]);
    if (!target_cloud || !input_cloud) return -1;

    std::cout << "Loaded " << target_cloud->size() << " data points from " << argv[1] << std::endl;
    std::cout << "Loaded " << input_cloud->size() << " data points from " << argv[2] << std::endl;

    // 体素滤波
    pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::ApproximateVoxelGrid<pcl::PointXYZ> approximate_voxel_filter;
    approximate_voxel_filter.setLeafSize(0.2, 0.2, 0.2);
    approximate_voxel_filter.setInputCloud(input_cloud);
    approximate_voxel_filter.filter(*filtered_cloud);
    std::cout << "Filtered cloud contains " << filtered_cloud->size() << " data points\n";

    // NDT配准
    pcl::NormalDistributionsTransform<pcl::PointXYZ, pcl::PointXYZ> ndt;
    ndt.setTransformationEpsilon(0.05);//0.01
    ndt.setStepSize(0.01);//0.05
    ndt.setResolution(0.1);//0.5
    ndt.setMaximumIterations(100);//50
    ndt.setInputSource(filtered_cloud);
    ndt.setInputTarget(target_cloud);

    Eigen::AngleAxisf init_rotation(0.6931, Eigen::Vector3f::UnitZ());
    Eigen::Translation3f init_translation(1.79387, 0.720047, 0);
    Eigen::Matrix4f init_guess = (init_translation * init_rotation).matrix();

    pcl::PointCloud<pcl::PointXYZ>::Ptr output_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    ndt.align(*output_cloud, init_guess);

    std::cout << "NDT converged: " << ndt.hasConverged() 
              << ", score: " << ndt.getFitnessScore() << std::endl;

    // 保存并可视化
    pcl::transformPointCloud(*input_cloud, *output_cloud, ndt.getFinalTransformation());
    pcl::io::savePCDFileASCII("output.pcd", *output_cloud);
    visualizer(target_cloud, output_cloud);

    return 0;
}
